import math as m
import numpy as np
norm = np.linalg.norm

# def main():
#     fin=open("1-forward.lammpstrj")
#     fout=open("a.xyz", "w+")
#     natom = 22
#     xyz = np.zeros([natom, 3])
#     vxyz = np.zeros([natom, 3])
#     atype = np.zeros(natom)
#     readxyz(fin, atype, xyz, vxyz)

def read_energy(name):
    e = np.loadtxt("{}.energy".format(name))
    return e

def read_intc(name):
    col = np.loadtxt("{}.intc".format(name))
    return col[1:, 1:]

def read_trj(name):
    try:
        with open("{}.lammpstrj".format(name)) as fin:
            fin.readline()
            fin.readline()
            fin.readline()
            a = fin.readline()
            natom = int(a.split()[0])
    except:
        raise NameError("the file cannot be read")

    xyz = np.zeros([natom, 3])
    vxyz = np.zeros([natom, 3])
    atype = np.zeros(natom)
    allxyz = []
    allvxyz = []
    with open("{}.lammpstrj".format(name)) as fin:
        while (readxyz(fin, atype, xyz, vxyz)):
            allxyz += [np.array(xyz.flat)]
            allvxyz += [np.array(vxyz.flat)]

    current_xyz = np.array(allxyz)
    current_vxyz = np.array(allvxyz)

    return current_xyz[1:], current_vxyz[1:]

def readxyz(fin, atype, xyz, vxyz):
    a = fin.readline()
    if (a):
      fin.readline()
      fin.readline()
      a = fin.readline()
      natom = int(a.split()[0])
      fin.readline()
      fin.readline()
      fin.readline()
      fin.readline()
      fin.readline()
      for i in range(natom):
          line = fin.readline().split()
          index = int(line[0])-1
          xyz[index, :] = np.array(list(map(float, line[2:5])))
          vxyz[index, :] = np.array(list(map(float, line[5:])))
          atype[index] = line[1]
      # xyz = rotateCtop(natom, xyz, 8, 10, 6)
      return True
    else:
        return False


def rotateCtop(natom, xyz0, centerid, topid, leftid):

    xyz = np.copy(xyz0)
    xyz = xyz - xyz[centerid, :]
    txyz = xyz[topid,:]
    dr = txyz
    r_tc = norm(dr)
    r_tcxy = norm(dr[:2])
    ct = dr[2]/r_tc
    st = m.sqrt(1-ct**2)
    cp = dr[0]/r_tcxy
    sp = dr[1]/r_tcxy
    rotation_matrix = np.array([
        [ct*cp, ct*sp, -st],
        [-sp, cp, 0],
        [st*cp, st*sp, ct]]).T
    # all xyz is operated by it
    xyz = xyz.dot(rotation_matrix)

    leftxyz = xyz[leftid,:]

    dr = leftxyz
    r_lcxy = norm(dr[:2])
    cp = dr[0]/r_lcxy
    sp = dr[1]/r_lcxy
    rotation_matrix = np.array([
        [cp, sp, 0],
        [-sp, cp, 0],
        [0, 0, 1.0]]).T
    # all xyz is operated by it
    xyz = xyz.dot(rotation_matrix)
    return xyz

# if __name__ == '__main__':
#     main()
